# 通过训练变分量子线路校正单量子比特信号（变分量子线路和神经网络同时训练，量子变分线路训练一步，神经网络训练一步，同步训练）

# 安装包包括：
# 1. qiskit
# 2. matplotlib
# 3. pylatexenc
# 4. seaborn
# 5. tensorflow
# 6. keras
# 7. qiskit_algorithms

import math
import matplotlib.pyplot as plt
from qiskit.visualization import plot_histogram
import numpy as np
import keras
from qiskit import QuantumCircuit, transpile
from qiskit.circuit import ParameterVector
from qiskit.quantum_info import Pauli, state_fidelity
from keras.models import Sequential
from keras.layers import Dense
from qiskit.primitives import Estimator
from qiskit_algorithms.gradients import ParamShiftEstimatorGradient

# 产生随机转动值（仿真对qubit的干扰噪音）
#random_rotations = np.random.uniform(0, 2 * np.pi, 3)
random_rotations = np.array([np.pi/2, 0, 0])
print("random_rotations=", random_rotations)
print("")

# 定义量子线路，其中rx ry rz为旋转门，遵从右手定则（例如对于rx，右手大拇指指向x轴方向，其他手指则为旋转方向）
qc = QuantumCircuit(1,1)
beta = ParameterVector('β', 3)
qc.rx(beta[0], 0)
qc.ry(beta[1], 0)
qc.rz(beta[2], 0)
qc.barrier()

# 定义受神经网络控制的量子线路，其中theta为控制的参数
theta = ParameterVector('θ', 3)
qc.rx(theta[0], 0)
qc.ry(theta[1], 0)
qc.rz(theta[2], 0)
qc.barrier()
# 拷贝一份电路加入测量，并没有实际作用，只是便于显示电路全貌。之所以不直接在原qc上加测量，是因为后面用到Estimator计算测量结果期望和梯度，电路中不能有测量操作
qc_measure = QuantumCircuit.copy(qc)
qc_measure.measure(0, 0)
# 显示量子线路
qc_measure.draw('mpl')
# 绑定随机旋转的3个角度参数
qc = qc.assign_parameters({beta: random_rotations})
qc_measure = qc_measure.assign_parameters({beta: random_rotations})
qc_measure.draw('mpl')

# 搭建神经网络模型
model = Sequential()
model.add(Dense(units=10, input_dim=1, activation='elu'))  # 1个输入command，units个输出
model.add(Dense(units=3)) #3个输出，控制θ1,θ2,θ3

# 编译模型
learning_rate = 0.05
optimizer = keras.optimizers.SGD(learning_rate=learning_rate)
model.compile(optimizer=optimizer, loss='mse')

# 初始化
theta_true = {0:[[0,0,0]],1:[[0,0,0]]}
theta_hat = {0:[[0,0,0]],1:[[0,0,0]]}
# 训练
k = 0
t = 0
while True:
    print("t=", t)
    # 命令
    command = np.random.choice([0, 1])
    #command = 0
    print("command=",command)
    # 预测
    theta_hat[command] = model.predict(np.array([command]), verbose=0)
    print("theta_hat=", theta_hat)

    # 计算量子线路测量的梯度和损失
    estimator = Estimator()
    expectation_hat = estimator.run(qc, Pauli('Z'), theta_hat[command]).result().values[0]
    expectation_true = 1 - command * 2  # 真实期望值（标记）
    print("expectation_hat, expectation_true=", expectation_hat, expectation_true)
    qc_loss = math.pow((expectation_hat - expectation_true) , 2)
    print("qc_loss=", qc_loss)
    if qc_loss < 0.001:
        break
    gradient = ParamShiftEstimatorGradient(estimator).run(qc, Pauli('Z'), theta_hat[command]).result().gradients[0]
    print('gradient =', gradient)
    theta_true[command] = theta_hat[command] - learning_rate * gradient * (expectation_hat - expectation_true)
    print('theta_true =', theta_true)
    nn_loss = model.train_on_batch(np.array([command]), theta_true[command])  # 神经网络损失
    print("nn_loss=", nn_loss)
    t = t + 1
print("")
print("训练完成：")
print("共训练", t, "轮")
print("qc_loss=", qc_loss)
print("theta_true=", theta_true)
print("")

print("最终测试结果：")
print("random_rotations=", random_rotations)
from qiskit import BasicAer
counts=[0,0]
for command in range(2):
    theta_hat[command] = model.predict(np.array([command]), verbose=0)
    backend = BasicAer.get_backend('qasm_simulator')
    bound_qc = qc.assign_parameters({theta: theta_hat[command][0]})
    bound_qc.measure(0, 0)
    # 输出计算结果
    job = backend.run(transpile(bound_qc, backend))
    counts[command] = job.result().get_counts(bound_qc)
    print("command=", command, ", counts=", counts[command])